import os
import os.path as osp
import sys
import torch
import torch.utils.data as data
import cv2
import numpy as np
import pandas as pd
import time
from PIL import Image, ImageDraw
import random
from .transform  import CenterNetDefaultTrainTransform

class CartoonFaceDetection(data.Dataset):
    def __init__(self, cfg, split='train'):
        '''
        cfg: config
        split: mode, 'train', 'val', 'test'
        preproc: transform 
        '''
        assert split in ('train', 'val', 'test'), 'dataset mode error'
        self.cfg = cfg
        self.split = split
        self.root_path = cfg.DATASET.ROOT
        self.anno_path = cfg.DATASET.ANNO_ROOT
        in_width, in_height = self.cfg.DATASET.IMAGE_SIZE

        pdcsv = pd.read_csv(self.anno_path)
        frames = {}
        with open(self.anno_path, 'r') as f:
            lines = f.readlines()
            for line in lines:
                line = line.split(',')
                im_name = line[0]
                x1, y1, x2, y2 = map(int, line[1:])
                if im_name not in frames:
                    frames[im_name]=[]
                frames[im_name].append([x1, y1, x2, y2])

        # for idx in range(len(pdcsv)):
        #     im_name = pdcsv.iloc[idx, 0]
        #     x1, y1, x2, y2 = pdcsv.iloc[idx, 1:]
        #     if im_name not in frames:
        #         frames[im_name]=[]
        #     frames[im_name].append([x1, y1, x2, y2])
        self.im_annos = list(frames.items())

        self.train_trans = CenterNetDefaultTrainTransform(in_width, in_height, 1, 4)
    
    def __getitem__(self, index):
        while True:
            img_path, annos = self.im_annos[index]
            # img = Image.open(osp.join(self.root_path, img_path))
            # img = img.convert('RGB')

            img = cv2.imread(osp.join(self.root_path, img_path))
            # # img = Image.fromarray(cv2.cvtColor(img,cv2.COLOR_BGR2RGB))
            # im_w, im_h = img.size
            im_h, im_w = img.shape[:2]
            ret = None
            targets = np.zeros((0, 5))
            for anno in annos:
                x1, y1, x2, y2 = anno[:4]
                x1, y1 = max(0, x1), max(0, y1)
                x2, y2 = min(x2, im_w), min(y2, im_h)
                label = 0 # face
                target = np.hstack(([x1, y1, x2, y2], label))
                targets = np.vstack((targets, target))
            
            ret = self.train_trans(img, targets)
            if ret['heatmap'].sum()>0:
                return ret
            else:
                index = random.randrange(0, len(self.im_annos))

    def __len__(self):
        return len(self.im_annos)


def detection_collate(batch):
    ret = {}
    for _, batch_i in enumerate(batch):
        for k, v in batch_i.items():
            if k not in ret:
                ret[k] = []
            ret[k].append(v)
    ret['img'] = np.stack(ret['img'], 0)
    return ret
